import pandas as pd
from sqlalchemy import create_engine
# 用于对分类变量进行标签编码的工具，将分类变量转换为整数形式，方便模型进行处理
from sklearn.preprocessing import LabelEncoder
# 用于划分数据集
from sklearn.model_selection import train_test_split
# 用于构建模型 -- 多元回归线性模型
from sklearn.linear_model import LinearRegression
import json

conn = create_engine('mysql+pymysql://root:123456@localhost:3306/house_data')
transfer = LabelEncoder()

# 该函数将其解析为JSON格式。如果解析成功，函数将返回解析后的JSON数据的最后一个元素（假设JSON数据是一个数组）；如果解析失败，则函数返回0。
def getJSONFirst(item):
    try:
        return int(json.loads(item)[-1])
    except:
        return 0

# 获取数据
def getData():
    df = pd.read_sql('select * from house_info', con=conn, index_col='id')
    X = df[['city', 'rooms_desc', 'area_range', 'houseType', 'sale_status', 'price']]
    X['houseType'] = X['houseType']\
        .replace('住宅', 0)\
        .replace('别墅', 1)\
        .replace('商业类', 2)\
        .replace('商业', 3)\
        .replace('酒店式公寓', 4)\
        .replace('底商', 5)\
        .replace('写字楼', 6)\
        .replace('车库', 7)
    # 通过transfer.fit_transform方法可以对城市数据进行标签编码，然后在预测时使用transfer.transform将新的城市数据转换为相同的编码形式
    X['city'] = transfer.fit_transform(X['city'].values)
    X['rooms_desc'] = X['rooms_desc'].apply(getJSONFirst)
    X['area_range'] = X['area_range'].apply(getJSONFirst)
    X['price'] = X['price'].astype('int')
    return X

# 模型训练
def model_train(data):
    # 数据集 测试集划分
    x_train,x_test,y_train,y_test =  train_test_split(data[['city','rooms_desc','area_range','houseType','sale_status']],data['price'],test_size=0.25,random_state=1)
    # 创建一个线性回归对象
    linear2 = LinearRegression()

    # 多元线性回归 模型训练
    linear2.fit(x_train,y_train)

    return linear2

# 预测
def pred(model,*args):
    city = transfer.transform([args[0]])[0]
    rooms_desc = args[1]
    area_range = args[2]
    houseType = ['住宅','别墅','商业类','酒店式公寓','底商','写字楼','车库'].index(args[3])
    sale_status = args[4]

    predPrice = model.predict([
        [city, rooms_desc, area_range, houseType, sale_status]
    ])
    return round(predPrice[0],1)



# 主函数
if __name__ == '__main__':
    model = model_train(getData())
    pred(model,'北京',5,150,'别墅',3)